Please address the comments from this code review:

## Overall Comments
- In `MiCoMLIRGen.convert`, the function argument is declared as `%input`, but `handle_placeholder` registers SSA names as `'%{n.name}'` (e.g. `%x`), so the first graph node never maps to `%input`; consider either naming the placeholder SSA `%input` or wiring the function argument SSA into `ssa_mapping` so the generated MLIR uses a consistent and defined input value.
- Several MLIR op emitters (e.g. `_handle_flatten`, pooling handlers, relu variants) use `result_type` for both operand and result types in the op signature, which is incorrect for shape-changing ops and brittle in general; it would be more robust to compute and pass the actual input tensor type separately (e.g. via the tracked `out` or predecessor node tensor) and only use `result_type` for the result.

## Individual Comments

### Comment 1
<location path="MiCoMLIRGen.py" line_range="255-264" />
<code_context>
+        bias_name = f"{n.name}_bias"
+        
+        self._add_weight_constant(weight_name, module.weight)
+        if module.bias is not None:
+            self._add_weight_constant(bias_name, module.bias)
+        
+        input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type
+        op = f"{result_ssa} = mico.linear({input_ssa}, @{weight_name}, @{bias_name}) : {input_type} -> {result_type}"
+        self._add_mlir_op(op)
+    
</code_context>
<issue_to_address>
**issue (bug_risk):** Guard against None bias being referenced in the generated MLIR

In these handlers, `bias_name` is always referenced (`@{bias_name}`) in the MLIR op, but the constant is only emitted when `module.bias is not None`. For modules created with `bias=False`, this produces a reference to an undeclared symbol. Please either skip the bias operand when `bias` is `None` or consistently emit a zero-initialized bias constant in that case.
</issue_to_address>

### Comment 2
<location path="MiCoMLIRGen.py" line_range="260-261" />
<code_context>
+        
+        # Generate operation
+        attrs = f"weight_bits = {module.qtype} : i32, act_bits = {module.act_q} : i32"
+        input_type = self._format_tensor_type(out) if hasattr(out, 'shape') else result_type
+        op = f"{result_ssa} = mico.bitlinear({input_ssa}, @{weight_name}, @{bias_name}) {{{attrs}}} : {input_type} -> {result_type}"
+        self._add_mlir_op(op)
+    
</code_context>
<issue_to_address>
**issue (bug_risk):** Input tensor type for ops is derived from the output tensor instead of the actual input

Across multiple handlers (`_handle_bitlinear`, `_handle_bitconv2d`, `_handle_bitconv1d`, `_handle_linear`, `_handle_conv2d`, `_handle_conv1d`), `input_type` is inferred from `out` or `result_type`, implicitly assuming input and output have the same shape/type. For conv/linear ops this is often false, so the generated MLIR will have incorrect operand types. Please derive `input_type` from the actual input tensor (via `self.node_info` or FX node args) instead of the result tensor.
</issue_to_address>

### Comment 3
<location path="MiCoMLIRGen.py" line_range="419" />
<code_context>
+    def _handle_flatten(self, n, module, input_ssa: str, result_ssa: str, result_type: str):
+        """Handle Flatten operation."""
+        start_dim = module.start_dim if hasattr(module, 'start_dim') else 1
+        op = f"{result_ssa} = mico.flatten({input_ssa}) {{start_dim = {start_dim} : i32}} : {result_type} -> {result_type}"
+        self._add_mlir_op(op)
+    
</code_context>
<issue_to_address>
**issue (bug_risk):** Flatten uses identical input/output types even though the shape usually changes

Both `_handle_flatten` and the functional `torch.flatten` path currently use the same `{result_type}` for input and output, even though flatten typically changes shape. This makes the IR types inaccurate and can break downstream passes. Please derive the input type from the pre-flatten tensor, the output type from `out`, and encode this as an explicit `input_type -> result_type` signature in both the module and function handlers.
</issue_to_address>

### Comment 4
<location path="MiCoMLIRGen.py" line_range="460-464" />
<code_context>
+            input_ssa = self._get_ssa(input_names[0])
+            op = f"{result_ssa} = mico.relu6({input_ssa}) : {result_type} -> {result_type}"
+            self._add_mlir_op(op)
+        elif function in (torch.flatten,):
+            input_ssa = self._get_ssa(input_names[0])
+            start_dim = input_args[1] if len(input_args) > 1 else 1
+            op = f"{result_ssa} = mico.flatten({input_ssa}) {{start_dim = {start_dim} : i32}} : {result_type}"
+            self._add_mlir_op(op)
+        elif function in (torch.cat,):
</code_context>
<issue_to_address>
**suggestion:** Flatten functional handler uses a different MLIR signature style than other ops

Here you only append `: {result_type}` without an `->` or explicit input type, while other ops use `: input_type -> result_type`. Please align this with the existing pattern (including both input and output types) so the generated IR remains consistent and matches the expected `mico.flatten` signature.

```suggestion
        elif function in (torch.flatten,):
            input_ssa = self._get_ssa(input_names[0])
            start_dim = input_args[1] if len(input_args) > 1 else 1
            op = (
                f"{result_ssa} = mico.flatten({input_ssa}) "
                f"{{start_dim = {start_dim} : i32}} : {result_type} -> {result_type}"
            )
            self._add_mlir_op(op)
```
</issue_to_address>

### Comment 5
<location path="MiCoTorchMLIRGen.py" line_range="161-169" />
<code_context>
+        
+        # Create a wrapper model that doesn't have MiCo-specific quantization
+        # torch-mlir expects standard PyTorch operations
+        class ModelWrapper(nn.Module):
+            def __init__(self, original_model):
+                super().__init__()
+                self.model = original_model
+            
+            def forward(self, x):
+                return self.model(x)
+        
+        wrapped_model = ModelWrapper(self.model)
+        wrapped_model.eval()
+        
</code_context>
<issue_to_address>
**question (bug_risk):** ModelWrapper doesn’t actually strip MiCo-specific layers before torch-mlir compilation

`ModelWrapper` just forwards to `self.model`, so it still exposes `BitLinear`/`BitConv*` and other MiCo-specific modules. If torch-mlir can’t handle these, this path may fail or yield invalid IR. Either adapt the wrapped model to present only standard `nn` layers (e.g., de-quantized / pre-lowered version) or update the comment to state that torch-mlir is expected to support MiCo QLayers directly.
</issue_to_address>

### Comment 6
<location path="tests/test_mlir_codegen.py" line_range="346" />
<code_context>
+        # Check for pooling operations (LeNet uses avgpool)
+        self.assertIn("mico.avgpool2d", content)
+    
+    def test_lenet_weight_constants(self):
+        """Test that LeNet generates weight constants."""
+        model = LeNet(1)
</code_context>
<issue_to_address>
**suggestion (testing):** Strengthen weight-constant tests by checking quantization metadata such as `scale` attributes and sub-byte types.

Right now this only verifies that weight constants exist in the MLIR, not that their quantization metadata is correct. Please extend the assertions to also cover:

- `scale = ... : f32` attributes on quantized weight constants for BitQLayers with `qw_scale`.
- Sub-byte tensor types like `!mico.int<4>`, `!mico.int<6>`, or `!mico.int<8>` where relevant (e.g., via `assertIn("!mico.int<4>", content)` or a `scale =` pattern).

This will ensure the serialized IR reflects the intended sub-byte quantization configuration, not just the presence of constants.
</issue_to_address>

### Comment 7
<location path="tests/test_torch_mlir_codegen.py" line_range="92" />
<code_context>
+        self.assertEqual(mlir_gen.is_torch_mlir_available(), TORCH_MLIR_AVAILABLE)
+
+
+class TestMiCoTorchMLIRGenConvert(unittest.TestCase):
+    """Test MLIR code conversion functionality."""
+    
</code_context>
<issue_to_address>
**suggestion (testing):** Introduce a test that exercises the torch-mlir path by stubbing `_convert_with_torch_mlir`, validating the combined template (Torch IR + MiCo overlay).

Current `TestMiCoTorchMLIRGenConvert` cases only cover fallback and generic file generation; they never drive the torch-mlir path. Even without torch-mlir installed, we can still cover this path by:

- Instantiating `MiCoTorchMLIRGen` and forcing `mlir_gen.use_torch_mlir = True`.
- Monkeypatching `_convert_with_torch_mlir` to return a dummy Torch dialect IR string.
- Calling `convert` and asserting that:
  - The output filename ends with `_torch_mlir.mlir`.
  - The file content includes both the dummy Torch IR and the MiCo overlay sections.

This adds coverage for the integration and template assembly logic without depending on a real torch-mlir installation.

Suggested implementation:

```python
        # Should match the global availability
        self.assertEqual(mlir_gen.is_torch_mlir_available(), TORCH_MLIR_AVAILABLE)


class TestMiCoTorchMLIRGenConvert(unittest.TestCase):
    """Test MLIR code conversion functionality."""

    def test_convert_uses_torch_mlir_template_and_mico_overlay(self):
        # Arrange: simple dummy model
        import torch

        class DummyModel(torch.nn.Module):
            def __init__(self):
                super().__init__()

            def forward(self, x):
                return x

        model = DummyModel()

        mlir_gen = MiCoTorchMLIRGen(model)
        # Force the torch-mlir path even if torch-mlir is not installed
        mlir_gen.use_torch_mlir = True

        dummy_torch_ir = "// DUMMY TORCH DIALECT IR\nmodule {}\n"

        # Stub out the actual torch-mlir conversion
        original_convert_with_torch_mlir = getattr(
            mlir_gen, "_convert_with_torch_mlir", None
        )

        def _fake_convert_with_torch_mlir(*args, **kwargs):
            return dummy_torch_ir

        try:
            setattr(mlir_gen, "_convert_with_torch_mlir", _fake_convert_with_torch_mlir)

            # Act: run conversion and capture output path
            with tempfile.TemporaryDirectory() as tmpdir:
                output_path = mlir_gen.convert(output_dir=tmpdir)

                # Assert: filename indicates torch-mlir path
                self.assertTrue(
                    output_path.endswith("_torch_mlir.mlir"),
                    msg=f"Expected torch-mlir output file, got: {output_path}",
                )

                # Assert: file content includes both the dummy Torch IR and MiCo overlay
                with open(output_path, "r", encoding="utf-8") as f:
                    content = f.read()

                self.assertIn(
                    dummy_torch_ir,
                    content,
                    msg="Torch dialect IR is missing from combined MLIR output.",
                )

                # The overlay marker / section name may differ in the real implementation;
                # this assertion is intentionally loose and should be adjusted if needed.
                self.assertRegex(
                    content,
                    r"MiCo|mico|MICO",
                    msg=(
                        "Expected MiCo overlay section to be present in combined MLIR "
                        "output, but no obvious MiCo markers were found."
                    ),
                )
        finally:
            # Restore original method if it existed
            if original_convert_with_torch_mlir is not None:
                setattr(
                    mlir_gen, "_convert_with_torch_mlir", original_convert_with_torch_mlir
                )



Tests the torch-mlir based MLIR code generation functionality including:
- Fallback behavior when torch-mlir is not available
- Basic initialization and configuration
- Output file generation
"""

import os
import sys
import tempfile
import shutil

```

1. The `convert` method signature for `MiCoTorchMLIRGen` is inferred as `convert(output_dir=...)`. If the real signature differs (e.g., `convert(self, path)` or returns `(path, content)`), update the call and assertions accordingly.
2. The filename suffix `_torch_mlir.mlir` is assumed from your comment. If the actual implementation uses a different suffix or naming convention, adjust the `endswith("_torch_mlir.mlir")` assertion to match.
3. The MiCo overlay assertion currently uses a loose regex for markers like `"MiCo"`, `"mico"`, or `"MICO"`. If the overlay has a specific, known marker (e.g., `"// MiCo overlay"` or a particular dialect name), replace the `assertRegex` with a more precise `assertIn` using that exact marker string.
4. Ensure that `MiCoTorchMLIRGen` is already imported in this test module. If not, add the appropriate import at the top of the file (e.g., `from mico.torch_mlir.codegen import MiCoTorchMLIRGen`) consistent with existing imports and project layout.
</issue_to_address>

### Comment 8
<location path="tests/test_torch_mlir_codegen.py" line_range="208-210" />
<code_context>
+        self.assertTrue(os.path.exists(mlir_path))
+
+
+class TestMiCoTorchMLIRGenReset(unittest.TestCase):
+    """Test reset functionality."""
+    
</code_context>
<issue_to_address>
**suggestion (testing):** Extend reset tests to verify that base-class MLIR state (ops, weights, SSA counters) is also cleared.

Right now the test only verifies the MiCoTorchMLIRGen-specific fields (`example_inputs`, `torch_mlir_ir`, `mico_quantization_ops`). Since `MiCoTorchMLIRGen.reset` delegates to `MiCoMLIRGen.reset`, please also assert that the base-class fields are reset (e.g. `mlir_ops` empty, `mlir_weights`/`mlir_weight_declarations` cleared, `ssa_counter` back to 0). You can populate them via a small forward pass, call `reset()`, and then check their values to guard against future regressions in the base reset logic.

```suggestion
class TestMiCoTorchMLIRGenReset(unittest.TestCase):
    """Test reset functionality."""

    def test_reset_clears_base_and_derived_state(self):
        # Build a simple model and run a forward pass to populate MLIR state.
        model = torch.nn.Sequential(
            torch.nn.Conv2d(1, 4, kernel_size=3),
            torch.nn.ReLU(),
        )
        mlir_gen = MiCoTorchMLIRGen(model)

        example_input = torch.randn(1, 1, 28, 28)
        mlir_gen.forward(example_input)

        # Sanity-check that both base-class and derived fields are populated
        # before we invoke reset(). This ensures the assertions after reset()
        # are meaningful.
        self.assertTrue(getattr(mlir_gen, "mlir_ops", []))
        self.assertTrue(getattr(mlir_gen, "mlir_weights", {}))
        self.assertTrue(getattr(mlir_gen, "mlir_weight_declarations", {}))
        self.assertNotEqual(getattr(mlir_gen, "ssa_counter", 0), 0)

        self.assertTrue(getattr(mlir_gen, "example_inputs", []))
        # torch_mlir_ir might be a string or IR object; just assert it's set.
        self.assertIsNotNone(getattr(mlir_gen, "torch_mlir_ir", None))
        self.assertTrue(getattr(mlir_gen, "mico_quantization_ops", []))

        # Now reset and verify that both base-class and MiCoTorchMLIRGen-specific
        # state has been cleared.
        mlir_gen.reset()

        # Base-class state
        self.assertEqual(len(getattr(mlir_gen, "mlir_ops", [])), 0)
        self.assertEqual(len(getattr(mlir_gen, "mlir_weights", {})), 0)
        self.assertEqual(len(getattr(mlir_gen, "mlir_weight_declarations", {})), 0)
        self.assertEqual(getattr(mlir_gen, "ssa_counter", 0), 0)

        # Derived-class state
        self.assertFalse(getattr(mlir_gen, "example_inputs", []))
        self.assertIsNone(getattr(mlir_gen, "torch_mlir_ir", None))
        self.assertEqual(len(getattr(mlir_gen, "mico_quantization_ops", [])), 0)
```
</issue_to_address>

### Comment 9
<location path="doc/MLIR_INTEGRATION.md" line_range="129" />
<code_context>
+```mlir
+// Sub-byte integer types
+!mico.int<1>   // 1-bit (binary)
+!mico.int<2>   // 2-bit (ternary)
+!mico.int<4>   // 4-bit
+!mico.int<8>   // 8-bit (standard)
</code_context>
<issue_to_address>
**issue (typo):** Clarify the comment for the 2-bit integer type, since "2-bit (ternary)" is inconsistent.

`// 2-bit (ternary)` is misleading: 2 bits allow 4 values, not 3. Please either change the description or adjust the bit width so the terminology matches.

```suggestion
!mico.int<2>   // 2-bit (4-level / quaternary)
```
</issue_to_address>